package com.rankeiot.ai.controller;

import com.rankeiot.ai.domain.Application;
import com.rankeiot.ai.domain.FunctionEntity;
import com.rankeiot.ai.message.ChatCompletionRequest;
import com.rankeiot.ai.message.ChatMessage;
import com.rankeiot.ai.message.Function;
import com.rankeiot.ai.message.FunctionParameters;
import com.rankeiot.ai.util.AiClient;
import com.rankeiot.ai.util.ErnieClient;
import com.rankeiot.ai.util.FunctionResult;
import com.rankeiot.ai.util.GptClient;
import com.rankeiot.core.Current;
import com.rankeiot.core.anno.Auth;
import com.rankeiot.core.data.UserInfo;
import com.rankeiot.core.data.response.Resp;
import com.rankeiot.core.util.JSONUtil;
import com.rankeiot.core.util.StringUtil;
import lombok.RequiredArgsConstructor;
import org.sagacity.sqltoy.dao.SqlToyLazyDao;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.messaging.handler.annotation.Header;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import org.ssssssss.magicapi.core.service.MagicAPIService;

import java.nio.charset.StandardCharsets;
import java.util.*;

@RequiredArgsConstructor
@Controller
@RequestMapping("ai")
public class AIController {
    private final SqlToyLazyDao dao;
    private final MagicAPIService magicAPIService;

    @Auth
    @PostMapping("/v1/chat/completions")
    public ResponseEntity<StreamingResponseBody> chatCompletion(@RequestHeader("Authorization") String authorization, @RequestBody ChatCompletionRequest body) {
        // 调用ChatGPT API进行聊天
        if (StringUtil.isEmpty(authorization)) {
            Resp.fail("密钥不正确");
        }
        String appKey = authorization.replaceFirst("Bearer ", "");
        Application application = new Application();
        application.setAppKey(appKey);
        application = dao.loadBySql("SELECT * FROM ai_application WHERE app_key=:appKey", application);
        if (application == null) {
            Resp.fail("密钥不正确");
        }
       AiClient realClient=null;
        if (application.getMode().startsWith("gpt")) {
            realClient = new GptClient(application);
        } else {
            realClient = new ErnieClient(application);
        }
        if (realClient == null) {
            Resp.fail("模型错误");
        }
        AiClient client = realClient;
        FunctionEntity functionQuery = new FunctionEntity();
        functionQuery.setAppId(application.getId());
        List<FunctionEntity> functions = dao.findBySql("SELECT * FROM ai_function WHERE app_id=:appId", functionQuery);

        Map<String,FunctionEntity> funMap=new HashMap<>();
        if (functions.size() > 0) {
            //加载前置提示词,设置函数调用
            Set<String> perScripts = new HashSet<>();
            // body.setFunctions();
            List<Function> funList = new ArrayList<>();
            for (FunctionEntity entity : functions) {
                String perScript = entity.getPerScript();
                if (StringUtil.isNotEmpty(perScript)) {
                    perScripts.add(entity.getPerScript());
                }
                funMap.put(entity.getId() + "", entity);
                Function fun = new Function();
                funList.add(fun);
                fun.setName(entity.getId() + "");
                fun.setDescription(entity.getDescription());
                if (StringUtil.isNotEmpty(entity.getParameters())) {
                    fun.setParameters(JSONUtil.parseObject(entity.getParameters(), FunctionParameters.class));
                }
                if (StringUtil.isNotEmpty(entity.getResponses())) {
                    fun.setResponses(JSONUtil.parseObject(entity.getParameters(), FunctionParameters.class));
                }
            }
            StringBuilder perPrompt = new StringBuilder();
            for (String perScript : perScripts) {
                Object ret = magicAPIService.invoke(perScript, new HashMap<>());
                if (ret != null) {
                    perPrompt.append(ret.toString());
                }
            }
            if (StringUtil.isNotEmpty(application.getPrompt())) {
                perPrompt.append(application.getPrompt());
            }
            body.setSystem(perPrompt.toString());

            body.setFunctions(funList);

        }
        UserInfo userInfo = Current.user();
        body.setUser_id(userInfo.getId());
        //获取历史信息
        StreamingResponseBody out = outputStream -> {
//            UserInfo userInfo = Current.user();
//            System.out.println(userInfo);
            System.out.println(JSONUtil.toJSONString(body));
            String outStr = client.chatCompletion(body,(funName,params)->{
                try {
                    FunctionEntity fun = funMap.get(funName);
                    Object ret = magicAPIService.invoke(fun.getScript(), params);
                    if(ret==null){
                        return new FunctionResult(fun.getDirect(), "");
                    }else if(ret instanceof String){
                        return new FunctionResult(fun.getDirect(), ret.toString());
                    }else if(ret instanceof FunctionResult){
                        FunctionResult fRet=(FunctionResult) ret;
                        if(fRet.getDirect()==null){
                            fRet.setDirect(fun.getDirect());
                            return fRet;
                        }
                    }
                    return new FunctionResult(fun.getDirect(), JSONUtil.toJSONString(ret));
                }catch (Exception e){
                    return new FunctionResult(true,e.getMessage());
                }
            });

            outputStream.write(outStr.getBytes(StandardCharsets.UTF_8));
            outputStream.flush();
            outputStream.close();
        };
        return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(out);
    }
}
